import os
from Tools import Constant
from Tools import VectorCollector
from sklearn.linear_model import LassoCV
import numpy as np
import json
import joblib



def solveVector(path,score):  
    filelist=os.listdir('Cache/'+path)
    vectors=[]
    print('score='+str(score))
    tot=0
    for x in filelist:
        name='Cache/'+path+x
        file=open(name,'r',encoding='UTF-8')
        paras=file.read()
        
        ss=paras.split('\n')
        paras=[]
        cnt=0
        for y in ss:
            if len(y)>0:
                paras.append(y)
                cnt+=1
        
        assert cnt==Constant.SentencesNumInPara
        try:
            vectors.append(VectorCollector.GetVector(paras))
        except:
            print('Error')
        file.close()
        tot+=1
        print(str(tot)+' '+str(x))
    storeVector(vectors,score)

def storeVector(vectors,score):

    X=vectors
    y=[]
    for x in range(0,len(X)):
        y.append(score)
    ans=[X,y]
    f=open('Cache/json/data'+str(score)+'.json','w',encoding='utf-8')
    json.dump(ans,f)    
    f.close()
'''
    f=open('Cache/json/data.json','r',encoding='utf-8')
    rd=json.load(f)
    f.close()
'''
#    trainVector(rd[0],rd[1])
def trainVector():
    filelist=os.listdir('Cache/json')
    vectors=[]
    scores=[]
    for js in filelist:
        f=open('Cache/json/'+js,'r',encoding='utf-8')
        rd=json.load(f)
        f.close()
        assert len(rd)==2
        vectors.extend(rd[0])
        scores.extend(rd[1])
    X=np.array(vectors)
    y=np.array(scores)
    reg = LassoCV(cv=5, random_state=0,positive=True).fit(X, y)
    joblib.dump(reg,'Cache/model.pkl')
def predictVector(s):
    reg=joblib.load('Cache/model.pkl')
    X=np.array(vectors)
    print(reg.predict(X[:]))
#solveVector(Constant.HighLevelSource,100)
#solveVector(Constant.MidLevelSource,50)
#solveVector(Constant.LowLevelSource,0) 

trainVector()
